"""
导入文本中的股票到mongodb股票池集合（stock_pool）
文本格式：
code \t name \t group1,group2...
文本位置：~/.astock/setting/stock_pool.txt
"""
import argparse
from AStock.ASSetting import settings, settings_path
import os
from pymongo import UpdateOne


def _main_(args):
    # file = args.file
    file = os.path.join(settings_path, 'stock_pool.txt')
    bulk = []
    with open(file, 'rt', encoding='utf-8') as f:
        for row in f:
            columns = row.split('\t')
            if len(columns) >= 2:
                code = columns[0].strip()
                name = columns[1].strip()
                groups = set(['全部'])
                if len(columns) >= 3:
                    grps = [grp.strip() for grp in columns[2].split(',')]
                    groups.update(grps)
                bulk.append(UpdateOne(
                    {'code': code, 'name': name},
                    {'$set': {'groups': list(groups)}},
                    upsert=True
                ))
                print('{} {} {}'.format(code, name, ','.join(groups)))
    if bulk:
        try:
            # 清空现有记录
            settings.database.stock_pool.drop()
            settings.database.stock_pool.bulk_write(bulk)
        except Exception as e:
            print(e)


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='import stocks to stock pool')
    # argparser.add_argument('-f', '--file', required=True, help='text file of stocks')

    args = argparser.parse_args()
    _main_(args)
